import matplotlib
from pydantic import BaseModel, Field

from langchain_core.tools import tool
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns


class PythonCodeInput(BaseModel):
    py_code: str = Field(description="一段合法的Python 代码字符串，例如'2+1'或 'x=3\\n y=x*2'")


@tool(args_schema=PythonCodeInput)
def python_inter(py_code: str) -> str:
    """
    当用户需要编写python程序并执行时,请调用该函数。
    该函数用于在当前Python环境中执行一段Python代码字符串，并返回执行结果。需要注意本函数只能执行非绘图的代码,
    若是绘图相关代码，则需要调用fig_inter函数运行
    """
    g = globals()
    try:
        return str(eval(py_code, g))
    except Exception as e:
        golobal_vars_before = set(g.keys())
        try:
            exec(py_code, g)
        except Exception as e:
            return f"执行代码报错{e}"
        golobal_vars_after = set(g.keys())
        new_vars = golobal_vars_after - golobal_vars_before
        if new_vars:
            result = {var: g[var] for var in new_vars}
            return str(result)
        else:
            return "代码执行成功"


class FigCodeInput(BaseModel):
    py_code: str = Field(description="要执行的Python绘图代码，必须使用matplotlib/seaborn创建图像并赋值给变量")
    fname: str = Field(description="图像对象的变量名，例如'fig'，用于从代码中提取并保存为图片")


@tool(args_schema=FigCodeInput)
def fig_inter(py_code: str,fname: str) -> str:
    """
    当用户需要编写python进行可视化绘图任务,请调用该函数。
    该函数会执行用户提供的Python绘图代码，并自动将生成的图像对象保存为图片文件并展示
    如果是要运行执行代码而不是绘图，请使用python_inter函数。
    """
    print("正在调用fig_inter工具运行python代码 ")
    current_backend = matplotlib.get_backend()
    matplotlib.use("Agg")

    local_vars = {"plt": plt, "pd": pd, "sns": sns}

    base_dir='/home/xnzn/disk/workspace/agent-chat-ui/public/images'
    try:
        g=globals()
        exec(py_code,g, local_vars)
        g.update(local_vars)

        fig=local_vars.get(fname,None)
        if fig is None:
            return f"绘图代码中未找到变量名为'{fname}'的图像对象，请检查代码是否正确"
        else:
            fig.savefig(f"{base_dir}/{fname}.png")
            return f"绘图代码执行成功，图像已保存为 {base_dir}/{fname}.png"
    except Exception as e:
        return f"执行绘图代码报错: {e}"